import numpy as np
from bert4keras.backend import keras,K
from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding,DataGenerator
from bert4keras.snippets import open
from bert4keras.layers import ConditionalRandomField
from keras.layers import Dense
from keras.models import Model
from tqdm import tqdm

maxlen = 256
epochs = 10
batch_size = 32
bert_layers = 12
learing_rate = 1e-5
crf_lr_multiplier = 1000


import os
import sys
if 'win' in sys.platform:
    BertPath = r'E:\nlp-data\pretrain_model\bert\model_chinese'
else:
    BertPath = r'/data/nlp_data/bert/model_chinese/'

config_path = os.path.join(BertPath,'bert_config.json')
checkpoint_path = os.path.join(BertPath,'bert_model.ckpt')
dict_path = os.path.join(BertPath,'vocab.txt')


def load_data(filename):
    D = []
    with open(filename,encoding='utf-8') as f:
        f = f.read()
        for l in f.split("\n\n"):
            if not l:
                continue
            d,last_flag = [],""
            for c in l.split("\n"):
                char,this_flag = c.split(" ")
                if this_flag == 'O' and last_flag == "O":
                    d[-1][0] += char
                elif this_flag == 'O' and last_flag != "O":
                    d.append([char,'O'])
                elif this_flag[:1] == 'B':
                    d.append([char,this_flag[2:]])
                else:
                    d[-1][0] += char
                last_flag = this_flag
            D.append(d)
    return D

train_data = load_data('datasets/china-people-daily-ner-corpus.tar/example.train')
valid_data = load_data('datasets/china-people-daily-ner-corpus.tar/example.dev')
test_data = load_data('datasets/china-people-daily-ner-corpus.tar/example.test')

tokenizer = Tokenizer(dict_path,do_lower_case=True)

# 类别映射
labels = ['PER', 'LOC', 'ORG']
id2label = dict(enumerate(labels))
label2id = {j:i for i,j in id2label.items()}
num_labels = len(labels) * 2  + 1

class data_generator(DataGenerator):
    def __iter__(self,random=False):
        batch_token_ids,batch_segment_ids,batch_labels = [],[],[]
        for is_end,item in self.sample(random):
            token_ids,labels = [tokenizer._token_start_id],[0]
            for w,l in item:
                w_token_ids = tokenizer.encode(w)[0][1:-1]
                if len(token_ids) + len(w_token_ids) < maxlen:
                    token_ids += w_token_ids
                    if l == 'O':
                        labels += [0] * len(w_token_ids)
                    else:
                        B = label2id[l] * 2 + 1
                        I = label2id[l] * 2 + 2
                        labels += ([B] + [I] * (len(w_token_ids) - 1))
                else:
                    break
            token_ids += [tokenizer._token_end_id]
            labels += [0]
            segment_ids = [0] * len(token_ids)
            batch_token_ids.append(token_ids)
            batch_segment_ids.append(segment_ids)
            batch_labels.append(labels)

            if len(batch_token_ids) == self.batch_size or is_end:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids,batch_segment_ids],batch_labels
                batch_token_ids,batch_segment_ids,batch_labels = [],[],[]

model = build_transformer_model(
    config_path,
    checkpoint_path
)

output_layer = 'Transformer-%s-FeedForward-Norm' % (bert_layers - 1)
output = model.get_layer(output_layer).output
output = Dense(num_labels)(output)

CRF = ConditionalRandomField(lr_multiplier=crf_lr_multiplier)
output  = CRF(output)

model = Model(model.input,output)
model.summary()

model.compile(
    loss = CRF.sparse_loss,
    optimizer=Adam(learing_rate),
    metrics = [CRF.sparse_accuracy]
)

def viterbi_decode(nodes,trans):
    """Viterbi算法求最优路径
        其中nodes.shape=[seq_len, num_labels],
            trans.shape=[num_labels, num_labels].
        """
    labels = np.arange(num_labels).reshape((1,-1))
    scores = nodes[0].reshape((-1,1))
    scores[1:] -= np.inf
    paths = labels
    for l in range(1,len(nodes)):
        M = scores + trans + nodes[l].reshape((1,-1))
        idxs = M.argmax(0)
        scores = M.max(0).reshape((-1,1))
        paths = np.concatenate([paths[:,idxs],labels],0)
    return paths[:,scores[:,0].argmax()]

def named_entity_recognize(text):
    """命名实体识别函数
        """
    tokens = tokenizer.tokenize(text)
    while len(tokens) > 512:
        tokens.pop(-2)
    mapping = tokenizer.rematch(text,tokens)
    token_ids = tokenizer.tokens_to_ids(tokens)
    segment_ids = [0] * len(token_ids)
    nodes = model.predict([[token_ids],[segment_ids]])[0]
    trans = K.eval(CRF.trans)
    labels = viterbi_decode(nodes,trans)
    entities,starting = [],False
    for i,label in enumerate(labels):
        if label > 0:
            if label % 2 == 1:
                starting = True
                entities.append([[i],id2label[(label - 1)//2]])
            elif starting:
                entities[-1][0].append(i)
            else:
                starting = False
        else:
            starting = False

    return [(
        text[mapping[w[0]][0]:mapping[w[-1]][-1] + 1],l) for w,l in entities
    ]

def evaluate(data):
    """评测函数
        """
    X,Y,Z = 1e-10,1e-10,1e-10
    for d in tqdm(data):
        text = ''.join([i[0] for i in d])
        R = set(named_entity_recognize(text))
        T = set([tuple(i) for i in d if i[1] != "O"])
        X += len(R &T)
        Y += len(R)
        Z += len(T)

    f1,precision,recall = 2 * X /( Y +Z),X/Y,X/Z
    return f1,precision,recall


class Evaluate(keras.callbacks.Callback):
    def __init__(self):
        self.best_val_f1 = 0

    def on_epoch_end(self, epoch, logs=None):
        trans = K.eval(CRF.trans)
        print(trans)
        f1,precision,recall = evaluate(valid_data)
        if f1 >= self.best_val_f1:
            self.best_val_f1 = f1
            model.save_weights('output/ner_crf.weights')
        print(
            'valid:  f1: %.5f, precision: %.5f, recall: %.5f, best f1: %.5f\n' %
            (f1, precision, recall, self.best_val_f1)
        )
        f1, precision, recall = evaluate(test_data)
        print(
            'test:  f1: %.5f, precision: %.5f, recall: %.5f\n' %
            (f1, precision, recall)
        )


if __name__ == '__main__':
    if not os.path.exists("output/ner_crf.weights"):
        evaluator = Evaluate()
        train_generator = data_generator(train_data,batch_size)

        model.fit_generator(
            train_generator.forfit(),
            steps_per_epoch=len(train_generator),
            epochs=epochs,
            callbacks=[evaluator]
        )

    else:
        train_generator = data_generator(train_data, batch_size)
        for i in train_generator.forfit():
            print(i)
            break
        model.load_weights("output/ner_crf.weights")
        evaluate(test_data)





